import csv


def train_txt_to_csv():
    data = []
    with open('../dataset/train.txt', 'r') as f:
        for line in f.readlines():
            data_line = line.split('\t')
            if data_line != '\t':
                data_line[-1] = data_line[-1].strip()
            data.append(data_line)
    data = deal_nan_data_train(data)
    with open('../dataset/train.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerows(data)


def test_txt_to_csv():
    data = []
    with open('../dataset/test.txt', 'r') as f:
        for line in f.readlines():
            data_line = line.split('\t')
            if data_line != '\t':
                data_line[-1] = data_line[-1].strip()
            data.append(data_line)
    data = deal_nan_data_test(data)
    with open('../dataset/test.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerows(data)


# 利用字典映射把字符串映射到数字上
def deal_vocab():
    data_train = []
    with open('../dataset/train.csv', 'r') as f:
        for line in f.readlines():
            data_line = line.split(',')
            if data_line != '\t':
                data_line[-1] = data_line[-1].strip()
            data_train.append(data_line)
    data_test = []
    with open('../dataset/test.csv', 'r') as f:
        for line in f.readlines():
            data_line = line.split(',')
            if data_line != '\t':
                data_line[-1] = data_line[-1].strip()
            data_test.append(data_line)

    dict_list = {}
    dict_num = 0
    for data_i in data_train:
        for j in range(26):
            if data_i[j+14] is '':
                data_i[j+14] = 'nan'
            if data_i[j+14] not in dict_list.keys():
                dict_list[data_i[j+14]] = dict_num
                dict_num += 1
    for data_i in data_test:
        for j in range(26):
            if data_i[j+13] is '':
                data_i[j+13] = 'nan'
            if data_i[j+13] not in dict_list.keys():
                dict_list[data_i[j+13]] = dict_num
                dict_num += 1
    for i in range(len(data_train)):
        for j in range(26):
            data_train[i][j+14] = str(dict_list[data_train[i][j+14]])
    for i in range(len(data_test)):
        for j in range(26):
            data_test[i][j+13] = str(dict_list[data_test[i][j+13]])
    with open('../dataset/train.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerows(data_train)
    with open('../dataset/test.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerows(data_test)


# 均值填充空值
def deal_nan_data_train(data):
    sum_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    count_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    mean_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    for i in data:
        for j in range(1, 14):
            if i[j] is not '':
                sum_[j] = sum_[j] + float(i[j])
                count_[j] += 1
    for j in range(1, 14):
        if count_[j] == 0:
            mean_[j] = 0
        else:
            mean_[j] = sum_[j] / count_[j]
    for i in data:
        for j in range(1, 14):
            if i[j] is '':
                i[j] = mean_[j]
    return data


def deal_nan_data_test(data):
    sum_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    count_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    mean_ = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    for i in data:
        for j in range(13):
            if i[j] is not '':
                sum_[j] = sum_[j] + float(i[j])
                count_[j] += 1
    for j in range(13):
        if count_[j] == 0:
            mean_[j] = 0
        else:
            mean_[j] = sum_[j] / count_[j]
    for i in data:
        for j in range(13):
            if i[j] is '':
                i[j] = mean_[j]
    return data


def read():
    train_txt_to_csv()
    test_txt_to_csv()
    deal_vocab()
